#include "ipu_protonet_classifier.hpp"
#include <iostream>
#include <fstream>
#include <filesystem>
#include <nlohmann/json.hpp>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <iomanip>

namespace protonet {

IPUProtoNetClassifier::IPUProtoNetClassifier(const std::string& model_path, 
                                            const std::pair<int, int>& input_size,
                                            const IPUImagePreprocessor::PreprocessConfig* preprocess_config)
    : input_size_(input_size) {
    
    std::cout << "🚀 正在初始化IPU ProtoNet分类器..." << std::endl;
    
    // 初始化IPU模型包装器
    model_wrapper_ = std::make_unique<IPUModelWrapper>(model_path, input_size);
    if (!model_wrapper_->isLoaded()) {
        throw std::runtime_error("IPU模型加载失败");
    }
    
    // 更新输入尺寸为模型的实际输入尺寸
    input_size_ = model_wrapper_->getInputSize();
    
    // 初始化图像预处理器
    image_preprocessor_ = std::make_unique<IPUImagePreprocessor>(input_size_, preprocess_config);
    
    std::cout << "✅ IPU ProtoNet分类器初始化成功" << std::endl;
    std::cout << "📁 模型路径: " << model_path << std::endl;
    std::cout << "📐 输入尺寸: " << input_size_.first << "x" << input_size_.second << std::endl;
}

bool IPUProtoNetClassifier::buildPrototypes(const std::string& dataset_path) {
    std::cout << "🔨 开始构建原型，数据集路径: " << dataset_path << std::endl;
    
    prototypes_.clear();
    class_names_.clear();
    
    try {
        // 遍历数据集目录
        for (const auto& class_entry : std::filesystem::directory_iterator(dataset_path)) {
            if (!class_entry.is_directory()) {
                continue;
            }
            
            std::string class_name = class_entry.path().filename().string();
            class_names_.push_back(class_name);
            std::cout << "\n📂 处理类别: " << class_name << std::endl;
            
            // 收集该类别的所有特征
            std::vector<std::vector<float>> class_features;
            
            for (const auto& image_entry : std::filesystem::directory_iterator(class_entry.path())) {
                if (!image_entry.is_regular_file()) {
                    continue;
                }
                
                std::string image_path = image_entry.path().string();
                std::cout << "  📸 处理图像: " << image_entry.path().filename().string() << std::endl;
                
                try {
                    // 提取特征
                    std::vector<float> features = extractFeatures(image_path);
                    if (!features.empty()) {
                        class_features.push_back(features);
                    }
                } catch (const std::exception& e) {
                    std::cerr << "    ⚠️ 处理图像失败: " << e.what() << std::endl;
                }
            }
            
            if (class_features.empty()) {
                std::cerr << "  ❌ 类别 " << class_name << " 没有有效的特征，跳过" << std::endl;
                class_names_.pop_back(); // 移除这个类别
                continue;
            }
            
            // 计算原型（均值特征）
            std::vector<float> prototype = computePrototype(class_features);
            prototypes_[class_name] = prototype;
            
            std::cout << "  ✅ 类别 " << class_name << " 原型构建完成 (样本数: " 
                      << class_features.size() << ", 特征维度: " << prototype.size() << ")" << std::endl;
        }
        
        std::cout << "\n🎯 原型构建完成！总共 " << prototypes_.size() << " 个类别" << std::endl;
        return !prototypes_.empty();
        
    } catch (const std::exception& e) {
        std::cerr << "❌ 构建原型时发生错误: " << e.what() << std::endl;
        return false;
    }
}

ClassificationResult IPUProtoNetClassifier::classify(const std::string& image_path,
                                                     DistanceMetric distance_metric,
                                                     ConfidenceMethod confidence_method) {
    try {
        if (prototypes_.empty()) {
            throw std::runtime_error("原型未构建，请先调用buildPrototypes()");
        }
        
        // 提取查询图像的特征
        std::vector<float> query_features = extractFeatures(image_path);
        
        // 计算与所有原型的距离
        std::vector<std::pair<std::string, float>> distances;
        for (const auto& [class_name, prototype] : prototypes_) {
            float distance = computeDistance(query_features, prototype, distance_metric);
            distances.emplace_back(class_name, distance);
        }
        
        // 根据距离度量排序（距离越小越相似）
        std::sort(distances.begin(), distances.end(),
                 [](const auto& a, const auto& b) { return a.second < b.second; });
        
        // 计算置信度
        float confidence = computeConfidence(distances, confidence_method);
        
        ClassificationResult result;
        result.predicted_class = distances[0].first;
        result.confidence = confidence;
        result.distance = distances[0].second;
        result.image_path = image_path;
        result.distance_metric = distance_metric;
        result.confidence_method = confidence_method;
        
        return result;
        
    } catch (const std::exception& e) {
        throw std::runtime_error("分类过程中发生错误: " + std::string(e.what()));
    }
}

std::vector<ClassificationResult> IPUProtoNetClassifier::evaluateDataset(
    const std::string& test_dataset_path,
    DistanceMetric distance_metric,
    ConfidenceMethod confidence_method) {
    
    std::cout << "🧪 开始评估数据集: " << test_dataset_path << std::endl;
    std::vector<ClassificationResult> results;
    
    try {
        for (const auto& class_entry : std::filesystem::directory_iterator(test_dataset_path)) {
            if (!class_entry.is_directory()) {
                continue;
            }
            
            std::string actual_class = class_entry.path().filename().string();
            std::cout << "\n📂 评估类别: " << actual_class << std::endl;
            
            for (const auto& image_entry : std::filesystem::directory_iterator(class_entry.path())) {
                if (!image_entry.is_regular_file()) {
                    continue;
                }
                
                std::string image_path = image_entry.path().string();
                std::cout << "  📸 " << image_entry.path().filename().string();
                
                try {
                    ClassificationResult result = classify(image_path, distance_metric, confidence_method);
                    result.actual_class = actual_class;
                    result.is_correct = (result.predicted_class == actual_class);
                    results.push_back(result);
                    
                    std::cout << " -> " << result.predicted_class 
                              << " (" << std::fixed << std::setprecision(3) << result.confidence << ")"
                              << " " << (result.is_correct ? "✅" : "❌") << std::endl;
                    
                } catch (const std::exception& e) {
                    std::cerr << " ❌ 错误: " << e.what() << std::endl;
                }
            }
        }
        
        // 计算准确率
        int correct_predictions = std::count_if(results.begin(), results.end(),
                                              [](const ClassificationResult& r) { return r.is_correct; });
        float accuracy = static_cast<float>(correct_predictions) / results.size();
        
        std::cout << "\n📊 评估结果:" << std::endl;
        std::cout << "   总样本数: " << results.size() << std::endl;
        std::cout << "   正确预测: " << correct_predictions << std::endl;
        std::cout << "   准确率: " << std::fixed << std::setprecision(2) 
                  << (accuracy * 100) << "%" << std::endl;
        
        return results;
        
    } catch (const std::exception& e) {
        throw std::runtime_error("评估数据集时发生错误: " + std::string(e.what()));
    }
}

std::vector<float> IPUProtoNetClassifier::extractFeatures(const std::string& image_path) {
    try {
        // 图像预处理
        std::vector<float> preprocessed_data = image_preprocessor_->preprocessImage(image_path);
        
        // IPU模型推理
        std::vector<float> features = model_wrapper_->infer(preprocessed_data);
        
        return features;
        
    } catch (const std::exception& e) {
        throw std::runtime_error("特征提取失败: " + std::string(e.what()));
    }
}

std::vector<float> IPUProtoNetClassifier::computePrototype(
    const std::vector<std::vector<float>>& features_list) {
    
    if (features_list.empty()) {
        return {};
    }
    
    size_t feature_dim = features_list[0].size();
    std::vector<float> prototype(feature_dim, 0.0f);
    
    // 计算平均值
    for (const auto& features : features_list) {
        for (size_t i = 0; i < feature_dim; ++i) {
            prototype[i] += features[i];
        }
    }
    
    // 归一化
    float scale = 1.0f / features_list.size();
    for (float& value : prototype) {
        value *= scale;
    }
    
    return prototype;
}

float IPUProtoNetClassifier::computeDistance(const std::vector<float>& features1,
                                           const std::vector<float>& features2,
                                           DistanceMetric metric) {
    if (features1.size() != features2.size()) {
        throw std::runtime_error("特征维度不匹配");
    }
    
    switch (metric) {
        case DistanceMetric::EUCLIDEAN: {
            float sum = 0.0f;
            for (size_t i = 0; i < features1.size(); ++i) {
                float diff = features1[i] - features2[i];
                sum += diff * diff;
            }
            return std::sqrt(sum);
        }
        
        case DistanceMetric::COSINE: {
            float dot_product = 0.0f, norm1 = 0.0f, norm2 = 0.0f;
            for (size_t i = 0; i < features1.size(); ++i) {
                dot_product += features1[i] * features2[i];
                norm1 += features1[i] * features1[i];
                norm2 += features2[i] * features2[i];
            }
            
            float cosine_sim = dot_product / (std::sqrt(norm1) * std::sqrt(norm2) + 1e-8f);
            return 1.0f - cosine_sim; // 转换为距离
        }
        
        default:
            throw std::runtime_error("不支持的距离度量");
    }
}

float IPUProtoNetClassifier::computeConfidence(
    const std::vector<std::pair<std::string, float>>& distances,
    ConfidenceMethod method) {
    
    if (distances.empty()) {
        return 0.0f;
    }
    
    switch (method) {
        case ConfidenceMethod::SIMPLE: {
            // 简单方法：基于最小距离的反比
            float min_distance = distances[0].second;
            return 1.0f / (1.0f + min_distance);
        }
        
        case ConfidenceMethod::NEGATIVE_LOG_LIKELIHOOD: {
            // 负对数似然方法
            std::vector<float> scores;
            for (const auto& [_, distance] : distances) {
                scores.push_back(-distance); // 距离越小，分数越高
            }
            
            // Softmax归一化
            float max_score = *std::max_element(scores.begin(), scores.end());
            float sum_exp = 0.0f;
            for (float& score : scores) {
                score = std::exp(score - max_score); // 减去最大值防止溢出
                sum_exp += score;
            }
            
            return scores[0] / sum_exp; // 返回最高分数类别的概率
        }
        
        default:
            return 0.0f;
    }
}

bool IPUProtoNetClassifier::savePrototypes(const std::string& file_path) const {
    try {
        nlohmann::json j;
        j["class_names"] = class_names_;
        j["input_size"] = {input_size_.first, input_size_.second};
        j["feature_dimension"] = model_wrapper_->getFeatureDimension();
        
        for (const auto& [class_name, prototype] : prototypes_) {
            j["prototypes"][class_name] = prototype;
        }
        
        std::ofstream file(file_path);
        file << j.dump(2);
        file.close();
        
        std::cout << "💾 原型已保存至: " << file_path << std::endl;
        return true;
        
    } catch (const std::exception& e) {
        std::cerr << "❌ 保存原型失败: " << e.what() << std::endl;
        return false;
    }
}

bool IPUProtoNetClassifier::loadPrototypes(const std::string& file_path) {
    try {
        std::ifstream file(file_path);
        if (!file.is_open()) {
            throw std::runtime_error("无法打开文件: " + file_path);
        }
        
        nlohmann::json j;
        file >> j;
        file.close();
        
        class_names_ = j["class_names"].get<std::vector<std::string>>();
        
        prototypes_.clear();
        for (const auto& [class_name, prototype_json] : j["prototypes"].items()) {
            prototypes_[class_name] = prototype_json.get<std::vector<float>>();
        }
        
        std::cout << "📂 原型已从文件加载: " << file_path << std::endl;
        std::cout << "   类别数: " << prototypes_.size() << std::endl;
        std::cout << "   特征维度: " << j["feature_dimension"].get<size_t>() << std::endl;
        
        return true;
        
    } catch (const std::exception& e) {
        std::cerr << "❌ 加载原型失败: " << e.what() << std::endl;
        return false;
    }
}

} // namespace protonet